"""
Transformer_CrossAsset Strategy Implementation
Uses Transformer model to capture cross-asset correlations and optimize portfolio weights
"""

from zipline.api import order_target_percent, symbol, schedule_function, get_datetime
import numpy as np
import pandas as pd

def initialize(context):
    """Initialize strategy parameters and schedules"""
    # Set Transformer parameters
    context.window = 90
    context.n_heads = 4
    context.d_model = 64
    context.n_layers = 3
    
    # Initialize mock Transformer model
    context.model = MockTransformerModel()
    
    # Asset symbols
    context.assets = {
        'SPY': symbol('SPY'),  # S&P 500 ETF
        'TLT': symbol('TLT'),  # 20+ Year Treasury Bond ETF
        'GLD': symbol('GLD'),  # Gold ETF
        'USO': symbol('USO')   # Crude Oil ETF
    }
    
    # Initialize embeddings and weights
    context.embeddings = {}
    context.target_weights = {asset: 0.25 for asset in context.assets}
    
    # Initialize asset returns tracking
    context.asset_returns = {}
    
    # Set scheduler for daily feature updates (market close)
    schedule_function(
        update_features,
        date_rule='every_day',
        time_rule='market_close'
    )
    
    # Set scheduler for weekly weight adjustments (Friday market close)
    schedule_function(
        adjust_weights,
        date_rule='week_end',
        time_rule='market_close'
    )
    
    # Set scheduler for monthly rebalancing (first day of month)
    schedule_function(
        rebalance,
        date_rule='month_start',
        time_rule='market_open'
    )

def handle_data(context, data):
    """Daily monitoring of asset performance"""
    # Track asset returns for volatility calculation
    for asset in context.assets.values():
        if data.can_trade(asset):
            price = data.current(asset, 'price')
            context.asset_returns[asset] = context.asset_returns.get(asset, []) + [price]

def update_features(context, data):
    """Update asset feature embeddings daily"""
    # Get asset returns data
    asset_returns = {}
    for sym, asset in context.assets.items():
        returns = context.asset_returns.get(asset, [])
        if len(returns) > 1:
            # Calculate daily returns
            daily_returns = [returns[i]/returns[i-1] - 1 for i in range(1, len(returns))]
            asset_returns[sym] = daily_returns[-context.window:]  # Last window days
    
    # Update embeddings using Transformer model
    if asset_returns:
        context.embeddings = context.model.update_embeddings(asset_returns)

def adjust_weights(context, data):
    """Adjust target weights based on attention and volatility"""
    if not context.embeddings:
        return
    
    # 1. Calculate attention weights
    attention_weights = calculate_attention(context.embeddings)
    
    # 2. Apply volatility adjustment
    context.target_weights = apply_volatility_adjustment(context, attention_weights)
    
    print("Adjusted weights:", context.target_weights)

def rebalance(context, data):
    """Rebalance portfolio to target weights"""
    # Apply single asset class limit (max 40%)
    for asset in context.target_weights:
        context.target_weights[asset] = min(context.target_weights[asset], 0.4)
    
    # Execute orders
    for asset, weight in context.target_weights.items():
        order_target_percent(context.assets[asset], weight)
    
    print("Portfolio rebalanced to target weights")

def calculate_attention(embeddings):
    """Calculate attention weights from embeddings"""
    # Mock implementation - real version would use Transformer attention
    weights = {}
    total = 0
    for asset in embeddings:
        # Simple weighting based on embedding magnitude
        weight = np.linalg.norm(embeddings[asset])
        weights[asset] = weight
        total += weight
    
    # Normalize weights
    return {asset: weight/total for asset, weight in weights.items()}

def apply_volatility_adjustment(context, weights):
    """Adjust weights based on volatility risk controls"""
    adjusted_weights = weights.copy()
    
    # Calculate volatilities
    volatilities = {}
    for asset in context.assets:
        returns = context.asset_returns.get(context.assets[asset], [])
        if len(returns) > 10:
            # Calculate 10-day volatility
            daily_returns = [returns[i]/returns[i-1] - 1 for i in range(1, len(returns))]
            volatilities[asset] = np.std(daily_returns[-10:])
    
    # Risk control: Reduce commodity weights if volatility > 30%
    if 'GLD' in weights and volatilities.get('GLD', 0) > 0.3:
        adjusted_weights['GLD'] *= 0.5
    if 'USO' in weights and volatilities.get('USO', 0) > 0.3:
        adjusted_weights['USO'] *= 0.5
    
    # Normalize weights
    total = sum(adjusted_weights.values())
    return {asset: weight/total for asset, weight in adjusted_weights.items()}

class MockTransformerModel:
    """Mock Transformer model for simulation purposes"""
    def __init__(self):
        self.embeddings = {}
    
    def update_embeddings(self, asset_returns):
        """Simulate embedding updates"""
        # Real implementation would use actual Transformer model
        # For simulation, generate random embeddings
        for asset, returns in asset_returns.items():
            # Create embedding based on return statistics
            mean_return = np.mean(returns) if returns else 0
            std_return = np.std(returns) if len(returns) > 1 else 0
            self.embeddings[asset] = np.array([mean_return, std_return])
        return self.embeddings